package com.zack.util.nlp.tfidf;

import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import com.zack.util.collections.MapProcessing;

/**
* author: dejian
* 2017  Sep 28, 2017  10:47:46 AM  (could be changed)
* 
* tf-idf = tf * idf
tf: 一篇文档里每个词在文档出现的次数 / 这篇文档的总词数
idf:  整个corpus里，文件总数 / 包含这个词的文件数+1
 因此tfidf 是针对每一篇文章里的每一个词来看的， 不是针对整个corpus的， 是贴合每一个文档的。
**/
public class CalculateTFIDF {

    public static void main(String[] args) {
        System.out.println("Prepare...");
        String wrtToFile = "data/tfidf/result/tfidf.result";
        // 获取整个corpus 的所有词的信息，按照doc分组。
        String corpusFile = "data/tfidf/corpus.txt";
        Map<String, List<TFIDFBean>> docTFIDFBeans = calculateTfIDFFromFile(corpusFile);
        PrintTFIDFUtil.printTFIDFBeansIntoFile(docTFIDFBeans,wrtToFile);
        System.out.println("Finish.");
    }
    
    
    /**
     * doc-2        背带裤 1   1   0.14285714285714285 0.07772400633575366
       doc-2       短裤  1   1   0.14285714285714285 0.07772400633575366
       doc-2       糖果色 1   1   0.14285714285714285 0.07772400633575366
       doc-3       三十  1   1   0.2 0.10881360887005515
       doc-3       胖   1   1   0.2 0.10881360887005515
       doc-3       外套  1   1   0.2 0.10881360887005515
     * 传入一个corpus， 文章用"\u2003\u2002\u2001"隔开。注意， 文章必须是分好词的。
     * 把一个文档集的每一个词，按照doc封装到TFIDFBean中
     * @param corpusPath
     * @param wrtToFile
     * @return
     */
    public static Map<String, List<TFIDFBean>> calculateTfIDFFromFile(String corpusPath) {
        TFAllFromFile.tfAllFromCorpusFile(corpusPath);
        
        System.out.println("Finish calculate tfidf from:"+corpusPath);
        return calculateTfIDF();
    }
    
    /**
     * 传入一个corpus dir path ， 一个文章在一个txt里面， 注意， 文章必须是分好词的。
     *  把一个文档集的每一个词，按照doc封装到TFIDFBean中
     * @param corpusPath
     * @param wrtToFile
     * @return
     */
    public static Map<String, List<TFIDFBean>> calculateTfIDFFromDir(String corpusDir) {
        TFAllFromFile.tfAllFromDir(corpusDir);
        System.out.println("Finish calculate tfidf from:"+corpusDir);
        return calculateTfIDF();
    }
    
    
    private static  Map<String, List<TFIDFBean>> calculateTfIDF(){
        Map<String, List<TFIDFBean>> fileTfidfBeans = new LinkedHashMap<>();
        for(Map.Entry<String, Map<String, Integer>> eachFileMap : TFAllFromFile.docTermFreqMap.entrySet()) { // 循环每一篇文档。
            String fileName = eachFileMap.getKey();
            
            if (null == TFAllFromFile.docNameTermNumMap.get(fileName)) {
                continue;
            }
            int docTermNum = TFAllFromFile.docNameTermNumMap.get(fileName);
            List<TFIDFBean> tfidfBeans = new LinkedList<>();
            // 每一个doc内，按照词频排序
            Map<String, Integer> sortedMap = MapProcessing.sortMapByMapValue(eachFileMap.getValue(), false);
            for(Map.Entry<String, Integer> termFreqEntry : sortedMap.entrySet()) {
                String term = termFreqEntry.getKey();
                double tf = termFreqEntry.getValue()*1.0 / (docTermNum*1.0);
                double idf = getIDFOfTerm(TFAllFromFile.docTermFreqMap,term);
                double tfidf = tf*idf;
                TFIDFBean tfidfBean = new TFIDFBean();
                tfidfBean.setDocName(fileName);
                tfidfBean.setIdf(idf);
                tfidfBean.setTerm(term);
                tfidfBean.setTermFreqInCorpus(TFAllFromFile.corpusTermFreqMap.get(term));
                tfidfBean.setTermFreqInDoc(TFAllFromFile.docTermFreqMap.get(fileName).get(term));
                tfidfBean.setTf(tf);
                tfidfBean.setTfidf(tfidf);
                tfidfBeans.add(tfidfBean);
            }
            fileTfidfBeans.put(fileName, tfidfBeans);
        }
        return fileTfidfBeans;
    }
    
    
    public static double getIDFOfTerm(Map<String, Map<String, Integer>> docTermFreqMap,String term) {
        int cntContains = 0; // 包含term的文档总数
        for(Map.Entry<String, Map<String, Integer>> eachFileMap : docTermFreqMap.entrySet()) {
            if (eachFileMap.getValue().containsKey(term)) {
                cntContains++;
            }
        }
       return Math.log10((docTermFreqMap.size()*1.0) / ((cntContains)+1));
    }
}
